Thanks for your PR. Could you also provide some visual output examples and comparison with other backends we support (latency, memory overhead, etc.)?

Thanks for your PR. Could you also provide some visual output examples and comparison with other backends we support (latency, memory overhead, etc.)?

@sayakpaul Thanks for the review! I just pushed two follow-up commits:

• 894e289 — fix: align ring_anything with the ulysses_anything path (use shared gather_size_by_comm, handle ring_anything in unshard hook).
• 09b3995 — docs: add ### Ring Anything Attention section in distributed_inference.md with config example, gloo TIP, current limitations, and the benchmark below.

FLUX.1-dev benchmarks on 4 × RTX 4090 (48GB)

ring_anything matches ring within ~0.4% latency and identical peak memory at evenly-divisible shapes, while it is the only ring path that works at non-divisible shapes such as 1008×1008. Full table is in the docs.

Visual outputs

I'll attach side-by-side generated images (same prompt + seed: single-GPU baseline vs. ring_anything) in a follow-up comment shortly.

Here are the visual outputs across all backends (same prompt, same seed).

Outputs

Comparison of generated images across different context-parallel backends.

resolution	ulysses	ring	unified_balanced	ulysses_anything	ring_anything
1024×1024
1008×1008	failed	failed	failed

Reproduction

import os
import time
import argparse
from typing import Any, Dict, List, Optional, Tuple

import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig


def make_step_timing_callback() -> Tuple[Any, List[float]]:
    step_times: List[float] = []
    last_time: Optional[float] = None

    def callback(pipe, step: int, timestep: int, callback_kwargs: Dict[str, Any]):
        nonlocal last_time
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        now = time.perf_counter()
        if last_time is not None:
            step_times.append(now - last_time)
        last_time = now
        return callback_kwargs

    return callback, step_times


def measure_memory():
    """Measure current GPU memory usage"""
    allocated = torch.cuda.memory_allocated() / 1024**3  # GB
    reserved = torch.cuda.memory_reserved() / 1024**3  # GB
    max_allocated = torch.cuda.max_memory_allocated() / 1024**3  # GB
    return {
        "allocated_gb": allocated,
        "reserved_gb": reserved,
        "max_allocated_gb": max_allocated,
    }

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--cp-backend",
        type=str,
        choices=["ring", "ulysses", "unified", 'ulysses_anything', 'ring_anything'],
        default="ulysses",
        help="Context parallel backend to use.",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=1024,
        help="Height of the generated image.",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=1024,
        help="Width of the generated image.",
    )
    return parser.parse_args()


def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="cpu:gloo,cuda:nccl")
    rank = dist.get_rank()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)
    return device


def main():
    args = parse_args()

    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        os.environ.get(
            "FLUX_DIR",
            "black-forest-labs/FLUX.1-dev",
        ),
        torch_dtype=torch.bfloat16,
    ).to(device)

    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()

    # Always using it because `ring` doesn't support default. This helps ensure a fair comparison.
    pipeline.transformer.set_attention_backend("_native_cudnn")

    if args.cp_backend == "ring":
        cp_config = ContextParallelConfig(ring_degree=world_size)
    elif args.cp_backend == "unified":
        cp_config = ContextParallelConfig(ring_degree=world_size // 2, ulysses_degree=world_size // 2)
    elif args.cp_backend == "ulysses":
        cp_config = ContextParallelConfig(ulysses_degree=world_size)
    elif args.cp_backend == "ulysses_anything":
        cp_config = ContextParallelConfig(ulysses_degree=world_size, ulysses_anything=True)
    elif args.cp_backend == "ring_anything":
        cp_config = ContextParallelConfig(ring_degree=world_size, ring_anything=True)
    else:
        raise ValueError(f"Unsupported cp_backend: {args.cp_backend}")

    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    def run_pipe(pipeline, steps=50, callback=None):
        image = pipeline(
            prompt,
            guidance_scale=3.5,
            num_inference_steps=steps,
            generator=torch.Generator().manual_seed(42),
            height=args.height,
            width=args.width,
            callback_on_step_end=callback,
        ).images[0]
        return image

    # warmup
    _ = run_pipe(pipeline, steps=10)

    callback, step_times = make_step_timing_callback()

    if torch.cuda.is_available():
        torch.cuda.synchronize()
    start = time.perf_counter()
    image = run_pipe(pipeline, steps=50, callback=callback)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    end = time.perf_counter()

    if dist.get_rank() == 0:
        print(f"{args.cp_backend=}")
        save_path = f"output_{args.height}x{args.width}_{args.cp_backend}.png"
        image.save(save_path)
        print(f"Saved image to {save_path}, time taken: {end - start:.2f} seconds")
        total_time = sum(step_times)
        iters = len(step_times)
        time_per_iter = (total_time / iters) * 1000  # in milliseconds
        steps_per_sec = iters / total_time
        print(f"Time per Iter: {time_per_iter:.2f} ms, Steps per Sec: {steps_per_sec:.2f}")

        max_allocated_gb = measure_memory()["max_allocated_gb"]
        print(f"{max_allocated_gb=}")


    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

# Example usage:
# torchrun --nproc_per_node=4 check_ulysses_anything.py --cp-backend ring # success
# torchrun --nproc_per_node=4 check_ulysses_anything.py --cp-backend ulysses # success
# torchrun --nproc_per_node=4 check_ulysses_anything.py --cp-backend unified # success
# torchrun --nproc_per_node=4 check_ulysses_anything.py --cp-backend ulysses_anything # success
# torchrun --nproc_per_node=4 check_ulysses_anything.py --cp-backend ring_anything # success

# torchrun --nproc_per_node=4 check_ulysses_anything.py --height 1008 --width 1008 --cp-backend ring # failed
# torchrun --nproc_per_node=4 check_ulysses_anything.py --height 1008 --width 1008 --cp-backend ulysses # failed
# torchrun --nproc_per_node=4 check_ulysses_anything.py --height 1008 --width 1008 --cp-backend unified # failed
# torchrun --nproc_per_node=4 check_ulysses_anything.py --height 1008 --width 1008 --cp-backend ulysses_anything # success
# torchrun --nproc_per_node=4 check_ulysses_anything.py --height 1008 --width 1008 --cp-backend ring_anything # success

Thanks for the updates -- they are very comprehensive. It seems like ulysses_anything is still the best but users who might prefer ring would want to use ring_anything for non-divisible shapes? Is that the motivation here?

Thanks for the updates -- they are very comprehensive. It seems like ulysses_anything is still the best but users who might prefer ring would want to use ring_anything for non-divisible shapes? Is that the motivation here?

Thanks for the review!

You're spot on — the primary motivation is to provide a seamless path for users on a ring topology to handle non-divisible sequence lengths, similar to how ulysses_anything works.

While ulysses/ulysses_anything is often the go-to, I believe ring_anything is a valuable addition for a few key reasons:

1. Zero Head-count Constraints: Unlike Ulysses, Ring doesn't require num_heads to be divisible by the degree. For models with irregular or small head counts, this provides a much-needed "fallthrough" option.
2. Performance on PCIe/P2P: In my experiments with dual-GPU PCIe setups, P2P ring communication can be more efficient than Ulysses' all-to-all due to reduced bandwidth contention across the PCIe switch.
3. Future work: This PR focuses on ensuring correctness and supporting general shapes. I plan to follow up with another PR to introduce a proper P2P implementation to fully unlock the computation-communication overlap benefits.

So while ulysses/ulysses_anything remains the recommended choice for most, while ring_anything covers head-constrained or ring-first setups—focusing on correctness today, with clear headroom for P2P overlap in the next step.

Your reasoning is spot on! Let's include them in the docs too.

@claude could you do a review here?

Claude finished @sayakpaul's task in 2m 29s —— View job

Review of PR #13545: Ring Anything Attention

• Read repository guidelines (AGENTS.md, models.md, modular.md)
• Read changed files and understand the implementation
• Compare with existing ulysses_anything and TemplatedRingAttention patterns
• Post detailed review feedback

Overall this is a well-structured PR that cleanly mirrors the existing patterns. The core algorithm (pad → all-gather → slice → online-softmax merge) is correct, and the config validation + hook routing are consistent with ulysses_anything. A few items below.

Duplicate hook branches should be merged

src/diffusers/hooks/context_parallel.py — both _prepare_cp_input (lines 213–220) and post_forward (lines 246–253) now have two identical branches:

if self.parallel_config.ulysses_anything:
    return PartitionAnythingSharder.shard_anything(...)
elif self.parallel_config.ring_anything:
    return PartitionAnythingSharder.shard_anything(...)

Since the bodies are identical, these should be merged:

if self.parallel_config.ulysses_anything or self.parallel_config.ring_anything:
    return PartitionAnythingSharder.shard_anything(...)

Same for the unshard_anything path in post_forward. This reduces duplication and makes it clearer that both modes use the same shard/unshard strategy.

Fix this →

Incomplete docstring comment for ring_anything field

src/diffusers/models/_modeling_parallel.py:88:

# Whether to enable ring anything attention to support sequence lengths.

This should read "to support arbitrary sequence lengths" (matching the docstring at line 68 which already says "arbitrary sequence lengths"). Minor, but worth keeping the inline comment consistent.

Memory considerations worth documenting

The current implementation pads all local KV to s_max before the all-gather. When sequence lengths are highly uneven across ranks, this communicates and allocates extra zeros. This is fine for a correctness-first approach (and the benchmarks confirm negligible overhead for near-uniform splits), but it would be good to add a brief inline comment noting this trade-off, especially since the PR description mentions a follow-up P2P implementation.

Minor: validation could be more explicit about mutual exclusivity

The current constraints in __post_init__ (lines 117–126) implicitly prevent ulysses_anything and ring_anything from being enabled simultaneously (since ulysses_anything requires ring_degree == 1 while ring_anything requires ulysses_degree == 1). Consider adding an explicit check:

if self.ulysses_anything and self.ring_anything:
    raise ValueError("ulysses_anything and ring_anything cannot both be enabled.")

This makes the constraint immediately obvious to readers without requiring them to reason through the degree constraints. It's also more robust against future changes to the degree validation.

Overall assessment

The core attention logic is correct:

• Pad → all-gather → slice correctly reconstructs each rank's true KV before calling the attention op (zeros never enter attention).
• The online-softmax merging loop is identical to TemplatedRingAttention and handles the variable-length chunks correctly.
• Config validation prevents invalid combinations.
• The hook routing through PartitionAnythingSharder correctly handles non-uniform splits.
• Backward raising NotImplementedError is consistent with TemplatedUlyssesAnythingAttention.

The benchmarks are thorough and confirm that ring_anything matches ring at divisible shapes (~0.4% overhead) while being the only ring-based mode that works at non-divisible shapes. The motivation (head-count constraints, PCIe/P2P topology advantages) is compelling.

The main actionable items are: (1) merging the duplicate hook branches, and (2) adding the explicit mutual-exclusivity check. The rest are minor polish.

@sayakpaul Thanks! Addressed the feedback by merging duplicate hooks, adding the exclusivity check, and polishing comments. Pushed in the latest commit.